package org.springframework.ai.chat.client.advisor;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.tuniu.agents.advisor.AgentChatMemoryAdvisor;
import com.tuniu.agents.advisor.ChatEnvironmentMemory;
import com.tuniu.agents.common.AgentsConstants;
import com.tuniu.agents.common.IdGenerator;
import com.tuniu.agents.common.TraceEnum;
import com.tuniu.agents.common.TraceLogVo;
import com.tuniu.agents.common.enums.SqlOperationType;
import com.tuniu.agents.message.TraceLogger;
import com.tuniu.agents.util.ChatResponseUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.client.advisor.api.*;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.model.MessageAggregator;
import org.springframework.util.StringUtils;
import reactor.core.publisher.Flux;

import java.util.Map;
import java.util.function.Function;

import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY;


public class SimpleLoggerAdvisor implements CallAroundAdvisor, StreamAroundAdvisor {

    public static final Function<AdvisedRequest, String> DEFAULT_REQUEST_TO_STRING = request -> request.toString();

    public static final Function<ChatResponse, String> DEFAULT_RESPONSE_TO_STRING = response -> {
        try {
            ObjectMapper objectMapper = new ObjectMapper();
            return objectMapper.writeValueAsString(response);
        } catch (Exception e) {
            return response.toString();
        }
    };

    private static final Logger logger = LoggerFactory.getLogger(SimpleLoggerAdvisor.class);

    private final Function<AdvisedRequest, String> requestToString;

    private final Function<ChatResponse, String> responseToString;

    private int order;

    protected final ChatEnvironmentMemory chatEnvironmentMemory;

    protected final TraceLogger traceLogger;

    protected final IdGenerator idGenerator;

    public SimpleLoggerAdvisor() {
        this(DEFAULT_REQUEST_TO_STRING, DEFAULT_RESPONSE_TO_STRING, 0, null, null,null);
    }

    public SimpleLoggerAdvisor(int order) {
        this(DEFAULT_REQUEST_TO_STRING, DEFAULT_RESPONSE_TO_STRING, order, null, null,null);
    }

    public SimpleLoggerAdvisor(Function<AdvisedRequest, String> requestToString,
                               Function<ChatResponse, String> responseToString, int order) {
        this(requestToString, responseToString, order, null, null,null);
    }

    public SimpleLoggerAdvisor(ChatEnvironmentMemory chatEnvironmentMemory, TraceLogger traceLogger, IdGenerator idGenerator) {
        this(DEFAULT_REQUEST_TO_STRING, DEFAULT_RESPONSE_TO_STRING, 0, chatEnvironmentMemory, traceLogger,idGenerator);
    }

    public SimpleLoggerAdvisor(Function<AdvisedRequest, String> requestToString,
                               Function<ChatResponse, String> responseToString,
                               int order,
                               ChatEnvironmentMemory chatEnvironmentMemory,
                               TraceLogger traceLogger,
                               IdGenerator idGenerator) {
        this.requestToString = requestToString;
        this.responseToString = responseToString;
        this.order = order;
        this.chatEnvironmentMemory = chatEnvironmentMemory;
        this.traceLogger = traceLogger;
        this.idGenerator = idGenerator;
    }

    @Override
    public String getName() {
        return this.getClass().getSimpleName();
    }

    @Override
    public int getOrder() {
        return this.order;
    }

    private AdvisedRequest before(AdvisedRequest request) {
        logger.debug("request: {}", this.requestToString.apply(request));
        //Store input information
        saveTrace(request,null);
        return request;
    }

    private void observeAfter(AdvisedRequest advisedRequest,AdvisedResponse advisedResponse) {
        logger.debug("response: {}", this.responseToString.apply(advisedResponse.response()));
        saveTrace(advisedRequest,advisedResponse);
    }

    private void saveTrace(AdvisedRequest request,AdvisedResponse response){
        String conversationId = request.advisorParams().get(CHAT_MEMORY_CONVERSATION_ID_KEY).toString();
        Map<String, Object> map =  chatEnvironmentMemory.get(conversationId);
        TraceLogVo traceLogVo = null;
        if(response == null) {
            //Store input information
            traceLogVo = new TraceLogVo.Builder()
                    .traceId(map.getOrDefault(AgentsConstants.AgentsConversationEnvironmentKey.PARAM_TRACE_ID, "").toString())
                    .userId(map.get(AgentsConstants.AgentsConversationEnvironmentKey.CUSTOMER_ID) != null ? map.get(AgentsConstants.AgentsConversationEnvironmentKey.CUSTOMER_ID).toString() : null)
                    .conversationId(map.get(AgentsConstants.AgentsConversationEnvironmentKey.CONVERSATION_ID).toString())
                    .requestId(map.get(AgentsConstants.AgentsConversationEnvironmentKey.CURRENT_REQUEST_ID).toString())
                    .token(0)
                    .agentId(request.advisorParams().get(AgentChatMemoryAdvisor.CHAT_MEMORY_AGENT_ID_KEY).toString())
                    .type(TraceEnum.LLM_INPUT.getDescription())
                    .content(request.userText())
                    .operationType(SqlOperationType.ChatTraceLogSaveTrace.getDescription())

                    .build();
        }else {
            //Store output information
            traceLogVo = new TraceLogVo.Builder()
                    .traceId(map.getOrDefault(AgentsConstants.AgentsConversationEnvironmentKey.PARAM_TRACE_ID, "").toString())
                    .userId(map.get(AgentsConstants.AgentsConversationEnvironmentKey.CUSTOMER_ID) != null ? map.get(AgentsConstants.AgentsConversationEnvironmentKey.CUSTOMER_ID).toString() : null)
                    .conversationId(map.get(AgentsConstants.AgentsConversationEnvironmentKey.CONVERSATION_ID).toString())
                    .requestId(map.get(AgentsConstants.AgentsConversationEnvironmentKey.CURRENT_REQUEST_ID).toString())
                    .token(response.response().getMetadata().getUsage().getCompletionTokens())
                    .agentId(request.advisorParams().get(AgentChatMemoryAdvisor.CHAT_MEMORY_AGENT_ID_KEY).toString())
                    .type(TraceEnum.LLM_OUTPUT.getDescription())
                    .tool(StringUtils.hasText(ChatResponseUtils.getFirstToolName(response.response()))?ChatResponseUtils.getFirstToolName(response.response()):"")
                    .content(ChatResponseUtils.getFirstTextMessage(response.response()))
                    .operationType(SqlOperationType.ChatTraceLogSaveTrace.getDescription())
                    .build();
            TraceLogVo UpdateTraceLogVo = new TraceLogVo.Builder()
                    .traceId(map.getOrDefault(AgentsConstants.AgentsConversationEnvironmentKey.PARAM_TRACE_ID, "").toString())
                    .requestId(map.get(AgentsConstants.AgentsConversationEnvironmentKey.CURRENT_REQUEST_ID).toString())
                    .token(response.response().getMetadata().getUsage().getPromptTokens())
                    .agentId(request.advisorParams().get(AgentChatMemoryAdvisor.CHAT_MEMORY_AGENT_ID_KEY).toString())
                    .type(TraceEnum.LLM_INPUT.getDescription())
                    .operationType(SqlOperationType.ChatTraceLogUpdateRecentLLMTrace.getDescription())
                    .build();
            traceLogger.updateRecentLLMTrace(UpdateTraceLogVo);
        }
        //save trace log
        traceLogger.saveTrace(traceLogVo);
    }

    @Override
    public String toString() {
        return SimpleLoggerAdvisor.class.getSimpleName();
    }

    @Override
    public AdvisedResponse aroundCall(AdvisedRequest advisedRequest, CallAroundAdvisorChain chain) {

        advisedRequest = before(advisedRequest);

        AdvisedResponse advisedResponse = chain.nextAroundCall(advisedRequest);

        observeAfter(advisedRequest,advisedResponse);

        return advisedResponse;
    }

    @Override
    public Flux<AdvisedResponse> aroundStream(AdvisedRequest advisedRequest, StreamAroundAdvisorChain chain) {

        advisedRequest = before(advisedRequest);

        AdvisedRequest processedRequest = AdvisedRequest.from(advisedRequest).build();

        Flux<AdvisedResponse> advisedResponses = chain.nextAroundStream(advisedRequest);

        return new MessageAggregator().aggregateAdvisedResponse(advisedResponses, response -> this.observeAfter(processedRequest, response));
    }

}